package com.tianji.aigc.memory;

import cn.hutool.core.lang.Assert;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.tianji.aigc.entity.ChatMessageMysql;
import com.tianji.aigc.enums.AgentTypeEnum;
import com.tianji.aigc.service.ChatMessageMysqlService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;

import java.util.List;

/**
 * 基于mysql实现的ChatMemoryRepository
 */
public class MysqlChatMemoryRepository implements ChatMemoryRepository {

    @Resource
    private ChatMessageMysqlService chatMessageMysqlService;


    @Override
    public List<String> findConversationIds() {
        List<ChatMessageMysql> mysqlMessages = this.chatMessageMysqlService.list(new QueryWrapper<>());
        return mysqlMessages.stream()
                .map(ChatMessageMysql::getConversationId)
                .distinct()
                .toList();
    }


    @Override
    public List<Message> findByConversationId(String conversationId) {
        QueryWrapper<ChatMessageMysql> wrapper = new QueryWrapper<>();
        wrapper.eq("conversation_id", conversationId);
        List<ChatMessageMysql> mysqlMessages = this.chatMessageMysqlService.list(wrapper);
        // 转换为Message对象列表
        return mysqlMessages.stream()
                .map(mysqlMessage -> MessageUtil.toMessage(mysqlMessage.getContent()))
                .toList();
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        Assert.notEmpty(messages, "消息列表不能为空");
        // 消息转换
        List<ChatMessageMysql> mysqlMessages = messages.stream()
                .filter(message -> {
                    // 过滤掉 agentType 的消息
                    String text = message.getText();
                    AgentTypeEnum agentTypeEnum = AgentTypeEnum.agentNameOf(text);
                    return agentTypeEnum == null;
                })
                .map(message -> ChatMessageMysql.fromMessage(conversationId, message))
                .toList();
        // 保存数据时，会传入全部的消息数据，包括之前的数据，所以需要先删除之前的数据，再添加新的数据
        this.deleteByConversationId(conversationId);
        // 存入数据库
        this.chatMessageMysqlService.saveBatch(mysqlMessages);
    }

    /**
     * 根据 conversationId 删除旧数据
     */
    @Override
    public void deleteByConversationId(String conversationId) {
        QueryWrapper<ChatMessageMysql> wrapper = new QueryWrapper<>();
        wrapper.eq("conversation_id", conversationId);
        this.chatMessageMysqlService.remove(wrapper);
    }

}


